import os
import math
import random
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import gym
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from sklearn.metrics import accuracy_score
from collections import deque

# from torch.nn.modules.module import T



"""

该文件在原先modify的基础上修改model，期望能做到模拟买入卖出以及持有等操作

全文均有修改，请与modify进行对比

"""


# 检查电脑gpu是否可用
if torch.cuda.is_available():
    print("ok we will use the gpu")
    torch.cuda.set_device(0)


# 表明环境参数个数
class observation_space:
    def __init__(self, n):
        self.shape = (n,)


# 表明接下来的举措，例如认为股票是涨还是跌
class action_space:
    def __init__(self, n):
        self.n = n

    def seed(self, seed):
        pass

    def sample(self):
        return random.randint(-1, self.n - 2)


# 进行模型的搭建，该模型使用pytorch进行搭建
# 这个对模型进行修改，模型的输出将变为限定到a和b这个闭区间内的整数，可以为负数
# 输出为正时表示买的股票数量，输出为负时表示卖的股票数量，为0时表明持有
class CustomModel(nn.Module):
    def __init__(self, input_size, a_size, b_size, hidden_units, output_size):
        super(CustomModel, self).__init__()
        self.input_size = input_size[0] * input_size[1]
        self.layer1 = nn.Linear(self.input_size, hidden_units)
        self.layer2 = nn.Linear(hidden_units, hidden_units)
        self.dropout = nn.Dropout(p=0.3)
        self.output_layer = nn.Linear(hidden_units, output_size)
        self.relu = nn.ReLU()


    def custom_activation(self, x, a, b):
        # 使用 sigmoid 激活函数将输出限制在 [0, 1] 区间
        x = torch.sigmoid(x)
        # 将其转换到 [a, b] 区间
        x = a + (b - a) * x
        # 对输出进行整数化
        x = torch.round(x)
        # 确保输出在 [a, b] 区间
        x = torch.clamp(x, min=a, max=b)
        return x

    def forward(self, x, a, b):
        x = x.flatten(0)
        x = self.relu(self.layer1(x))
        x = self.dropout(x)
        x = self.relu(self.layer2(x))
        x = self.dropout(x)
        x = self.output_layer(x)

        # 使用自定义激活函数将输出限制在 [a, b] 区间
        x = self.custom_activation(x, a, b)

        return x


# 进行股票环境的搭建
class Finance:
    # 需要预测的股票路径
    # url = '000001.SZ.csv'

    # window，窗口大小，用于计算滚动统计量，如移动平均、收益率的移动平均和标准差。
    # leverage，杠杆倍数，用于放大收益或损失。它影响每个交易的回报计算。这里暂时不做改动
    # min_performance，最低绩效阈值，用于决定是否结束训练或评估过程。如果性能低于此阈值，训练或评估过程将终止。
    # start，数据的起始索引，表示从哪个数据点开始进行训练或评估。用于在数据中选择合适的时间范围。
    # end，数据的结束索引
    # mu，数据的均值，用于标准化数据。如果为 None，则在数据准备过程中计算并使用数据的均值。
    # std，数据的标准差，用于标准化数据。如果为 None，则在数据准备过程中计算并使用数据的标准差。
    def __init__(self, symbol, features, window, lags, leverage=1, min_performance=0.85, start=0, end=None, mu=None,
                 std=None):
        self.bef = 10000 # 记录上一周用户手里剩余资产（包含现金以及股票总价值）
        self.money = 10000 # 模拟用户手里有多少钱
        self.amount = 0 # 目前用户手里有多少支股票
        self.symbol = symbol # symbol，指定需要预测的股票的符号（即股票的代码），在这里也把他变成一个参数计算为特征值，暂不理解
        self.features = features # features，模型训练与预测时需要用到的参数，也是股票环境的参数，例如 features = [symbol, 'r', 's', 'm', 'v']
        self.n_features = len(features)
        self.window = window
        self.lags = lags # lags，回溯期，用于定义状态空间的长度。如果设置为 10，则表示状态空间将包括过去 10 个时间步的数据。
        self.start = start
        self.end = end
        self.mu = mu
        self.std = std
        self.observation_space = observation_space(self.lags)  # 观测空间，即环境的大小，也可以认为是要观察的天数的多少
        self.action_space = action_space(3)  # 动作空间，即可以进行的动作的类别数量
        self._get_data()
        self._prepare_data()
        self.now_values = self.data_['close'].iloc[0] # 当前股票价格
        self.aft_values = self.data_['aft_close'].iloc[0] # 明天股票价格

    def _get_data(self):
        dataframes = []
        dataframes.append(pd.read_csv('000001.SZ.csv', index_col=0, parse_dates=True).dropna())
        dataframes.append(pd.read_csv('000002.SZ.csv', index_col=0, parse_dates=True).dropna())
        dataframes.append(pd.read_csv('000004.SZ.csv', index_col=0, parse_dates=True).dropna())
        dataframes.append(pd.read_csv('000006.SZ.csv', index_col=0, parse_dates=True).dropna())
        dataframes.append(pd.read_csv('000007.SZ.csv', index_col=0, parse_dates=True).dropna())
        dataframes.append(pd.read_csv('000008.SZ.csv', index_col=0, parse_dates=True).dropna())
        dataframes.append(pd.read_csv('000009.SZ.csv', index_col=0, parse_dates=True).dropna())

        self.raw = pd.concat(dataframes)
        # self.raw = pd.read_csv('000001.SZ.csv', index_col=0, parse_dates=True).dropna()


    # 对于获得的数据进行预处理，这个是计算平均值，涨跌等参数，如果读取的数据已经是处理好了的，可以不使用该函数
    def _prepare_data(self):
        self.data = pd.DataFrame(self.raw)
        self.data.dropna(inplace=True)
        self.data['aft_close'] = self.data['close'].shift(-1) # 记录明天的收盘值
        self.data.dropna(inplace=True)
        self.data = self.data.iloc[self.start:]
        self.data.dropna(inplace=True)

        print("look hear")
        print(len(self.data))

        self.data_ = self.data.copy()

        if self.end is not None:
            self.data = self.data.iloc[:self.end - self.start]
            self.data_ = self.data_.iloc[:self.end - self.start]


    # 将连续lags的数据当做当前的环境参数，该参数将会丢入模型进行训练
    def _get_state(self):
        return self.data_[self.features].iloc[self.bar - self.lags:self.bar].values

    def seed(self, seed=None):
        random.seed(seed)
        np.random.seed(seed)

    # 重置环境，相当于一轮新的训练开始
    def reset(self):
        self.money = 10000
        self.bef = 10000
        self.amount = 0
        self.now_values = self.data_['close'].iloc[0] # 当前股票价格
        self.aft_values = self.data_['aft_close'].iloc[0] # 明天股票价格
        self.bar = self.lags
        state = self.data_[self.features].iloc[self.bar - self.lags:self.bar]
        return state.values

    # 对于当前环境进行预测之后会进入新的环境。例如预测了明天的涨跌之后，接下来更新环境开始预测后天的涨跌
    # 运行后返回新的环境参数
    def step(self, action):
        reward = 0 # 这里暂时不考虑如何加奖励

        self.amount += action
        self.money -= action * self.data['close'].iloc[self.bar]

        self.bar += 1
        if self.bar < len(self.data_):
            self.now_values = self.data['close'].iloc[self.bar]
            self.aft_values = self.data['aft_close'].iloc[self.bar]


        if self.bar >= len(self.data):
            done = True
        else:
            done = False

        state = self._get_state()
        info = {}
        return state, reward, done, info  # 返回值依次为环境参数state，奖励分数reward，是否结束当前学习done（注意done为True时预测结束或者说是失败），以及额外信息info


def set_seeds(seed=100):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


# 进行学习体的搭建
class FQLAgent(nn.Module):
    # gamma，折扣因子，gamma越接近1越注重未来奖励
    # hidden_units，隐藏层的单元数量
    # opt，优化器
    # learning_rate，学习率
    # learn_env，用于训练的环境
    # valid_env，用于训练的环境
    def __init__(self, gamma=0.95, hidden_units=24, opt=optim.RMSprop, learning_rate=0.001, learn_env=None,
                 valid_env=None):
        super(FQLAgent, self).__init__()
        self.learn_env = learn_env
        self.valid_env = valid_env
        self.epsilon = 1.0  # 探索率，用于控制动作的随机性，探索率越大，模型越有概率随机进行探索
        self.epsilon_min = 0.01  # 最低探索率
        self.epsilon_decay = 0.995  # 探索率衰减因子，随着模型训练的增加，模型的探索率就会越变越小，更注重于自身训练的结果
        self.gamma = gamma
        self.batch_size = 5  # 样本批量
        self.max_treward = 0  # 有记录的最大累计奖励
        self.hidden_units = hidden_units
        self.trewards = list()  # 存储每个训练回合的累计奖励
        self.averages = list()  # 存储每个训练回合的奖励平均值，用于观察奖励的变化趋势
        self.performances = list()  # 存储每个训练回合的环境性能指标
        self.aperformances = list()  # 存储每 25 个回合的性能指标的滑动平均值，用于平滑性能波动。
        self.vperformances = list()  # 存储每个验证回合的环境性能指标。
        self.memory = deque()  # 经验回放池，用于存储代理的经验，以便后续进行学习，这里会进行未来奖励
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 检查CUDA是否可用并选择设备
        # 进行模型的搭建
        self.model, self.criterion, self.optimizer = self._build_model((self.learn_env.lags, self.learn_env.n_features),
                                                                       1, 1,
                                                                       self.hidden_units, 1, opt, learning_rate)
        # print("show lags and n_features: ",self.learn_env.lags, self.learn_env.n_features)
        self.model = self.model.to(self.device)  # 将模型或数据迁移到gpu上，加快运行

    def _build_model(self, input_size, a_size, b_size, hidden_units, output_size, optimizer, learning_rate):
        model = CustomModel(input_size, a_size, b_size, hidden_units, output_size)
        criterion = nn.MSELoss()  # 损失函数
        optimizer = optimizer(model.parameters(), lr=learning_rate)  # 优化器
        return model, criterion, optimizer

    def act(self, state, l, r):  # 这里进行一次行为的预测，即根据当前环境模型进行预测下次的行为
        if random.random() <= self.epsilon:  # 探索率越大，越有可能进行随机选择某种行为，适用于模型训练前期
            return random.randint(l, r)
        else:  # 模型进行预测
            state = torch.tensor(state, dtype=torch.float32).unsqueeze(0).to(self.device)  # 调整数据类型和形状，将输入数据移动到GPU上
            action = self.model(state, l, r)  # 得到结果
            return action.item()

    def replay(self):  # 进行经验回放，也就是将之前预测过的行为重新观看一下，站在现有的结果看待之前的预测是否成功，并进行未来奖励
        batch = random.sample(self.memory, self.batch_size)  # 随机抽取memory中一个批次的数据
        for today in batch: # 每个单位内存一周信息，七天
            for state, action, l, r, reward, next_state, done in today:
                state_tensor = torch.tensor(state, dtype=torch.float32).unsqueeze(0).to(self.device)
                next_state_tensor = torch.tensor(next_state, dtype=torch.float32).unsqueeze(0).to(self.device)
                # 按照上述修改后的结果开始训练模型
                target = self.model(state_tensor, l, r).detach()
                loss = F.mse_loss(self.model(state_tensor, l, r), target)
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
            if self.epsilon > self.epsilon_min:
                self.epsilon *= self.epsilon_decay

    # 进行学习体的学习
    def learn(self, episodes):  # 总共学习episodes轮
        actions = []
        for e in range(1, episodes + 1):
            av = 0
            state = self.learn_env.reset()  # 每次训练前重置模型
            # state = np.reshape(state, [1, self.learn_env.lags, self.learn_env.n_features])
            for _ in range(5000):
                l = -self.learn_env.amount
                r = math.floor(self.learn_env.money / self.learn_env.now_values)
                action = self.act(state, l, r)  # 预测当前环境下选什么行为比较好
                next_state, reward, done, info = self.learn_env.step(action)  # 执行预测后的行为后的新环境以及得到的奖励等
                state = next_state  # 更新当前环境，之后的学习预测要在新环境上进行
                actions.append([state , action, l, r, reward, next_state, done])

                if _ % 7 == 0: # 每隔一周观测营收情况
                    # 如果一周过去进行结算的时候发现赚钱了就丢入经验池中
                    if self.learn_env.money + self.learn_env.amount * self.learn_env.aft_values > self.learn_env.bef:
                        self.memory.append(actions)
                    actions.clear()
                    self.learn_env.bef = self.learn_env.money + self.learn_env.amount * self.learn_env.aft_values

                if done:  # 如果done=True，表示这一轮失败或者说是结束，进行数据的记录
                    treward = self.learn_env.money + self.learn_env.amount * self.learn_env.aft_values
                    self.max_treward = max(self.max_treward, treward)
                    self.trewards.append(treward)
                    print(f'episode: {e}/{episodes} | date : {_} | treward: {treward} | max: {self.max_treward}')
                    break

            # self.validate(e, episodes)  # 测试
            if len(self.memory) > self.batch_size:  # 当经验回放池中数据足够多时，进行经验回放
                self.replay()

    # 在验证环境中评估模型的性能
    def validate(self, e, episodes):
        state = self.valid_env.reset()
        state = torch.tensor(state, dtype=torch.float32).unsqueeze(0).to(self.device)  # 将输入数据移动到GPU上
        # print("state.shape: ",state.shape, "and type: ", )
        # state = np.reshape(state, [1, self.valid_env.lags, self.valid_env.n_features])
        for _ in range(10000):
            action_values = self.model(state)
            action = torch.argmax(action_values).item()
            next_state, reward, done, info = self.valid_env.step(action)
            # state = np.reshape(next_state, [1, self.valid_env.lags, self.valid_env.n_features])
            if done:
                treward = _ + 1
                # perf = self.valid_env.performance
                # self.vperformances.append(perf)
                if e % 20 == 0:
                    templ = 71 * '='
                    templ += '\nepisode: {:2d}/{} | VALIDATION | '
                    templ += 'treward: {:4d} | '
                    templ += 'eps: {:.2f}\n'
                    templ += 71 * '='
                    print(templ.format(e, episodes, treward, self.epsilon))
                break
    # def test(self, episodes):
    #     trewards = []
    #     for e in range(1, episodes + 1):
    #         state = env.reset()
    #         for _ in range(501):
    #             state_tensor = torch.tensor(state, dtype=torch.float32).unsqueeze(0).to(self.device)
    #             action = torch.argmax(self.model(state_tensor)).item()
    #             next_state, reward, done, info = env.step(action)
    #             state = next_state
    #             if done:
    #                 treward = _ + 1
    #                 trewards.append(treward)
    #                 print(f'episode: {e}/{episodes} | treward: {treward}')
    #                 break
    #     return trewards


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("looooooooooook this: ", device)

# 参数调整
symbol = '000001.SZ'
# features = ['close', 'open', 'high', 'low', 'vol']
features = ['close', 'open', 'high', 'low', 'pre_close', 'change', 'pct_chg', 'vol', 'amount', 'MACD', 'DEA',
            'Histogram', 'MA5', 'MA10', 'RSI', 'MOM', 'EMA12', 'EMA26', 'open-close', 'high-low']

a = 0
b = 5000
c = 0

# 训练环境搭建
learn_env = Finance(symbol, features, window=10, lags=6, start=1, end=a + b, mu=None, std=None)

# print(learn_env.data.info())

# 测试环境搭建
# valid_env = Finance(symbol, features, window=learn_env.window, lags=learn_env.lags,
#                     start=a + b, end=a + b + c, mu=learn_env.mu, std=learn_env.std)

set_seeds(100)
# 学习体搭建
agent = FQLAgent(learning_rate=0.0001, learn_env=learn_env, valid_env=None)

episodes = 1001

# 开始训练
agent.learn(episodes)
print(agent.epsilon)

# 绘图
plt.figure(figsize=(10, 6))
x = range(1, len(agent.trewards) + 1)
y = np.polyval(np.polyfit(x, agent.trewards, deg=3), x)
plt.plot(agent.trewards, label='moving money')
plt.plot(x, y, 'r--', label='trend')
plt.xlabel('episodes')
plt.ylabel('total money')
plt.legend()
plt.show()
